package com.katze.common.ws.core;

import com.katze.common.ws.config.WebSocketProperties;
import com.katze.common.ws.handler.HeartbeatHandler;
import com.katze.common.ws.handler.MessageDecoderHandler;
import com.katze.common.ws.handler.MessageHandler;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelPipeline;
import io.netty.channel.socket.SocketChannel;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslContextBuilder;
import io.netty.handler.ssl.util.InsecureTrustManagerFactory;
import io.netty.handler.timeout.IdleStateHandler;
import org.springframework.core.io.ClassPathResource;

import javax.net.ssl.KeyManagerFactory;
import java.io.InputStream;
import java.security.KeyStore;
import java.util.concurrent.TimeUnit;

public class DefaultChannelInitializer extends ChannelInitializer<SocketChannel> {

    private final MessageDispatcher dispatcher;
    private final WebSocketProperties properties;
    private final SslContext sslContext;

    public DefaultChannelInitializer(MessageDispatcher dispatcher, WebSocketProperties properties) {
        this.dispatcher = dispatcher;
        this.properties = properties;
        this.sslContext = createSslContext(properties.getSsl());
    }

    @Override
    protected void initChannel(SocketChannel ch) {
        ChannelPipeline pipeline = ch.pipeline();

        // SSL 支持
        if (sslContext != null) {
            pipeline.addLast(sslContext.newHandler(pipeline.channel().alloc()));
        }

        // HTTP协议编解码
        pipeline.addLast(new HttpServerCodec());
        pipeline.addLast(new HttpObjectAggregator(65536));

        // WebSocket协议处理器
        pipeline.addLast(new WebSocketServerProtocolHandler(properties.getPath(), null, true));

        // 心跳检测
        if (!properties.getIdleTimeout().isNegative()) {
            pipeline.addLast(new IdleStateHandler(properties.getIdleTimeout().toMillis(), 0, 0, TimeUnit.MILLISECONDS));
            pipeline.addLast(new HeartbeatHandler());
        }


        // 自定义协议处理
        pipeline.addLast(new MessageDecoderHandler());
        pipeline.addLast(new MessageHandler(dispatcher));
    }


    private SslContext createSslContext(WebSocketProperties.Ssl ssl) {
        if (ssl == null) {
            return null;
        }
        try (InputStream keyStoreStream = new ClassPathResource(ssl.getKeyStorePath()).getInputStream()) {
            KeyStore ks = KeyStore.getInstance("JKS");
            ks.load(keyStoreStream, ssl.getKeyStorePassword().toCharArray());
            KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            kmf.init(ks, ssl.getKeyStorePassword().toCharArray());

            return SslContextBuilder.forServer(kmf)
                    .trustManager(InsecureTrustManagerFactory.INSTANCE) // 生产环境应配置CA证书
                    .build();
        } catch (Exception e) {
            throw new IllegalStateException("Failed to load SSL context");
        }
    }
}
